# -*- coding: utf-8 -*-
"""
Created on Thu Apr 21 23:31:41 2022

@author: 13503
"""

import os
import torch
import numpy as np
import h5py
import dgl
#import dgl.nn.functional as dglfn
from dgl.data import DGLDataset

class SignNetGraph(DGLDataset):
    def __init__(self, 
                 graph_dataset_name:str,
                 dataset_path:list[str], 
                 feat_path:list[str])->None:
        self.dataset_path = dataset_path
        self.feat_path = feat_path
        
        dataset_dir, dataset_name = os.path.split(dataset_path)
        super().__init__(name=graph_dataset_name, raw_dir=dataset_dir, save_dir=dataset_dir)

    def process(self) -> None:
        # read data from file to generate Dataset
        print("enter process function")
        np_data = np.transpose(np.array(h5py.File(name=self.dataset_path)["Gwl_ud"])).astype(np.float16)
        data_in_tensor = torch.from_numpy(np_data).type(torch.float) # transform array into tensor
        src, dst = torch.nonzero(data_in_tensor, as_tuple=True)

        self.graph = dgl.graph((src,dst), num_nodes=data_in_tensor.shape[0])
        self.graph.edata["sign"] = data_in_tensor[src,dst].reshape((self.graph.num_edges(),1))
        self.graph.ndata["feat"] = torch.load(self.feat_path).cpu()

        pos_neg = torch.zeros((self.graph.num_nodes(),2)) # store how many positive and negative links starts from each node
        signs = torch.zeros((self.graph.num_edges(),1)) # store the sign of each link
        
        A_pos = A_neg = torch.zeros(data_in_tensor.shape)
        pos_src = neg_dst = []

        for i in range(len(src)):
            if data_in_tensor[src[i]][dst[i]] == 1:
                pos_neg[src[i]][0] += 1
                A_pos[src[i]][dst[i]] = 1
            else:
                pos_neg[src[i]][1] += 1
                A_neg[src[i]][dst[i]] = -1
        
        print("There are {0} edges in total ".format(self.graph.num_edges()))
        print("count the number of negative links = {0}".format(torch.sum(pos_neg.index_select(1,torch.tensor([1])))))
        print("count the number of positve links = {0}".format(torch.sum(pos_neg.index_select(1,torch.tensor([0])))))
        for i in range(len(src)):
            if data_in_tensor[src[i]][dst[i]] > 0:
                signs[i] = 1.0 / pos_neg[src[i]][0] 
            else:
                signs[i] = -10.0 / pos_neg[src[i]][1]

        self.graph.ndata["pos_neg"] = pos_neg
        self.graph.edata["sign"] = signs
        

    def __getitem__(self, idx)->None:
        # get one of the grahs and store it onto device
        return self.graph
    
    def __len__(self)-> None:
        return 1

    def save(self)->None:
        # save graphs onto disk
        dgl.save_graphs(self.graph_path, self.graph)

    def load(self)->None:
        # load processed data from directory `self.save_path`
        self.graph, _ = dgl.load_graphs(self.graph_path)
        self.graph = self.graph[0]

    def has_cache(self)->bool:
        # check whether there are processed data in `self.save_path`
        self.graph_path = os.path.join(self.save_dir, self.name.split(".")[0] + '_dgl_graph.bin')
        print("self.graph_path = {0}".format(self.graph_path))
        return os.path.exists(self.graph_path)
